import torch
from torch import nn
from copy import deepcopy
from .base import Attacker
from torch.cuda import amp
from .PGD import PGD

class PGD_cross(PGD):
    def __init__(self, model1, model2, img_transform=(lambda x:x, lambda x:x), use_amp=False):
        super().__init__(model1,  img_transform=(lambda x:x, lambda x:x), use_amp=False)
        self.model2 = model2

    def step(self, images, labels, loss):

        images.requires_grad = True
        outputs = self.model(images)
        outputs1 = self.model2(images)

        self.model.zero_grad()
        self.model2.zero_grad()
        cost = loss(outputs, labels)
        cost1 = loss(outputs1, labels)
        cost = cost + cost1
        cost.backward()

        adv_images = (self.img_transform[1](images) + self.alpha() * images.grad.sign()).detach_()
        eta = torch.clamp(adv_images - self.img_transform[1](self.ori_images), min=-self.eps, max=self.eps)
        images = self.img_transform[0](torch.clamp(self.img_transform[1](self.ori_images) + eta, min=0, max=1).detach_())

        return images

    def attack(self, images, labels):
        self.ori_images = deepcopy(images)

        for i in range(self.iters):
            self.model.eval()
            self.model2.eval()

            images = self.forward(self, images, labels)

            self.model.zero_grad()
            self.model2.zero_grad()
            self.model.train()
            self.model2.train()
        return images